This PR fixes a tensor model conversion bug that prevented PyTorch models with GeneralFitting types (dipole, polar) from being converted to TensorFlow format.

Problem

When converting PyTorch models with GeneralFitting types (dipole, polar) to TensorFlow using dp convert-backend, the conversion failed with:

ValueError: cannot reshape array of size 6 into shape (2,100)

This error occurred because the TensorFlow deserialization code attempted to reshape the out_bias tensor to match the bias_atom_e tensor shape, but these tensors have fundamentally different purposes and incompatible shapes in GeneralFitting models:

• out_bias: Shape varies by fitting type (e.g., [1, ntypes, 3] for dipole) - represents the output bias
• bias_atom_e: Shape [ntypes, embedding_width] (e.g., [2, 100]) - represents the internal fitting network bias

Solution

Modified the deserialization logic in deepmd/tf/model/model.py to:

1. Categorize fitting types correctly: Distinguish between InvarFitting types (ener, dos, property) and GeneralFitting types (dipole, polar)
2. Apply reshape logic selectively: Only apply the original reshape logic to InvarFitting types where out_bias and bias_atom_e have compatible shapes
3. Skip reshape for GeneralFitting types: Avoid the incompatible reshape operation for dipole and polar models
4. Preserve backward compatibility: Keep the original logic for energy models and other InvarFitting types

Fitting Types Handled

• InvarFitting types (reshape applied): ener, energy, dos, property, direct_force, direct_force_ener
• GeneralFitting types (reshape skipped): dipole, polar

Reproduction Steps

The issue can be reproduced with:

cd examples/water_tensor/dipole
dp --pt train dipole_input_torch.json --skip-neighbor-stat
dp --pt freeze -o dw_model.pth
dp convert-backend dw_model.pth dw_model.pb  # Previously failed, now works

cd ../polar
dp --pt train polar_input_torch.json --skip-neighbor-stat
dp --pt freeze -o pw_model.pth
dp convert-backend pw_model.pth pw_model.pb  # Also works now

Testing

• ✅ Successfully converts PyTorch dipole models to TensorFlow
• ✅ Successfully converts PyTorch polar models to TensorFlow
• ✅ Converted models load and function correctly
• ✅ Energy models continue to work (no regression)
• ✅ All fitting types handled correctly based on their class hierarchy

Fixes #4411.

💬 Share your feedback on Copilot coding agent for the chance to win a $200 gift card! Click here to start the survey.